"""
基于Transformer的时序预测策略
策略逻辑：
1. 使用Transformer模型预测未来价格走势
2. 结合注意力机制识别关键时间点
3. 动态调整交易频率和仓位
"""

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

class TimeSeriesDataset(Dataset):
    def __init__(self, prices, seq_length=60):
        self.prices = prices
        self.seq_length = seq_length
        
    def __len__(self):
        return len(self.prices) - self.seq_length
        
    def __getitem__(self, idx):
        seq = self.prices[idx:idx+self.seq_length]
        target = self.prices[idx+self.seq_length]
        return torch.FloatTensor(seq), torch.FloatTensor([target])

class TransformerModel(nn.Module):
    def __init__(self, input_dim=1, hidden_dim=64, n_heads=4, n_layers=3):
        super().__init__()
        self.encoder = nn.Linear(input_dim, hidden_dim)
        self.pos_encoder = PositionalEncoding(hidden_dim)
        encoder_layer = nn.TransformerEncoderLayer(hidden_dim, n_heads)
        self.transformer = nn.TransformerEncoder(encoder_layer, n_layers)
        self.decoder = nn.Linear(hidden_dim, 1)
        
    def forward(self, x):
        x = self.encoder(x)
        x = self.pos_encoder(x)
        x = self.transformer(x)
        x = self.decoder(x[-1])  # 只取最后一个时间步
        return x

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + self.pe[:x.size(0)]
        return x

class TransformerForecastStrategy:
    def __init__(self, ticker, seq_length=60):
        self.ticker = ticker
        self.seq_length = seq_length
        self.model = None
        
    def train(self, prices, epochs=50):
        """训练Transformer模型"""
        dataset = TimeSeriesDataset(prices, self.seq_length)
        dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
        
        self.model = TransformerModel()
        criterion = nn.MSELoss()
        optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
        
        for epoch in range(epochs):
            for seq, target in dataloader:
                optimizer.zero_grad()
                output = self.model(seq.unsqueeze(-1))
                loss = criterion(output, target)
                loss.backward()
                optimizer.step()
                
    def predict(self, recent_prices):
        """预测未来价格走势"""
        if self.model is None or len(recent_prices) < self.seq_length:
            return 0
            
        with torch.no_grad():
            seq = torch.FloatTensor(recent_prices[-self.seq_length:]).unsqueeze(-1)
            pred = self.model(seq).item()
            last_price = recent_prices[-1]
            return (pred - last_price) / last_price  # 预测收益率

if __name__ == '__main__':
    strategy = TransformerForecastStrategy('AAPL')
    
    # 模拟价格数据
    np.random.seed(42)
    prices = np.cumsum(np.random.randn(1000)) + 100
    
    # 训练模型
    strategy.train(prices)
    
    # 测试预测
    pred_return = strategy.predict(prices[-60:])
    print(f"预测收益率: {pred_return:.2%}")
    print("交易信号:", "买入" if pred_return > 0.01 else "卖出" if pred_return < -0.01 else "持有")